"""
黄金交易预测项目 - 主程序
整合数据获取、特征工程、模型训练和预测功能
"""

import os
import sys
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
import argparse
import warnings
warnings.filterwarnings('ignore')

# 导入项目模块
from src.get_data import generate_mock_gold_data, generate_mock_usd_data, save_data
from src.create_features import process_combined_data, save_features
from src.ml_models import GoldMLModels


class GoldTradingPredictor:
    """黄金交易预测器主类"""
    
    def __init__(self, data_dir="data", model_dir="models"):
        self.data_dir = data_dir
        self.model_dir = model_dir
        self.ml_models = GoldMLModels(model_dir)
        
        # 确保目录存在
        os.makedirs(os.path.join(data_dir, "raw"), exist_ok=True)
        os.makedirs(os.path.join(data_dir, "features"), exist_ok=True)
        os.makedirs(os.path.join(data_dir, "processed"), exist_ok=True)
        os.makedirs(model_dir, exist_ok=True)
    
    def run_full_pipeline(self, start_date="2014-01-01", end_date=None):
        """
        运行完整的机器学习流水线
        
        Args:
            start_date (str): 开始日期
            end_date (str): 结束日期
        """
        print("=" * 80)
        print("黄金交易预测项目 - 完整流水线")
        print("=" * 80)
        
        if end_date is None:
            end_date = datetime.now().strftime("%Y-%m-%d")
        
        print(f"数据时间范围: {start_date} 到 {end_date}")
        print()
        
        # 步骤1: 数据获取
        print("步骤1: 数据获取")
        print("-" * 40)
        self._acquire_data(start_date, end_date)
        print()
        
        # 步骤2: 特征工程
        print("步骤2: 特征工程")
        print("-" * 40)
        self._engineer_features()
        print()
        
        # 步骤3: 模型训练
        print("步骤3: 模型训练")
        print("-" * 40)
        self._train_models()
        print()
        
        # 步骤4: 模型评估
        print("步骤4: 模型评估")
        print("-" * 40)
        self._evaluate_models()
        print()
        
        print("完整流水线执行完成！")
        print("=" * 80)
    
    def _acquire_data(self, start_date, end_date):
        """数据获取步骤"""
        print("生成模拟数据...")
        
        # 生成黄金数据
        gold_data = generate_mock_gold_data(start_date, end_date)
        save_data(gold_data, 'gold_spot_data.csv')
        
        # 生成美元指数数据
        usd_data = generate_mock_usd_data(start_date, end_date)
        save_data(usd_data, 'usd_index_data.csv')
        
        # 合并数据
        combined_data = pd.merge(
            gold_data[['date', 'close']].rename(columns={'close': 'gold_close'}),
            usd_data[['date', 'close']].rename(columns={'close': 'usd_close'}),
            on='date',
            how='inner'
        )
        
        save_data(combined_data, 'combined_data.csv')
        print(f"数据获取完成: {len(combined_data)} 条记录")
    
    def _engineer_features(self):
        """特征工程步骤"""
        # 加载合并数据
        combined_file = os.path.join(self.data_dir, "raw", "combined_data.csv")
        combined_data = pd.read_csv(combined_file)
        combined_data['date'] = pd.to_datetime(combined_data['date'])
        
        # 创建特征
        features = process_combined_data(combined_data)
        
        # 保存特征
        save_features(features, 'gold_features.csv')
        print(f"特征工程完成: {features.shape} 数据，{len(features.columns)} 个特征")
    
    def _train_models(self):
        """模型训练步骤"""
        # 加载特征数据
        features_file = os.path.join(self.data_dir, "features", "gold_features.csv")
        features_data = pd.read_csv(features_file)
        
        # 训练回归模型
        print("训练回归模型...")
        X_train, X_test, y_train, y_test, feature_cols = self.ml_models.prepare_data(
            features_data, target_col='target_return_1d', test_size=0.2
        )
        
        regression_models = self.ml_models.train_regression_models(
            X_train, y_train, model_names=['rf', 'xgb', 'lgb', 'ridge']
        )
        
        # 训练分类模型
        print("训练分类模型...")
        X_train_cls, X_test_cls, y_train_cls, y_test_cls, feature_cols_cls = self.ml_models.prepare_data(
            features_data, target_col='target_direction_1d', test_size=0.2
        )
        
        classification_models = self.ml_models.train_classification_models(
            X_train_cls, y_train_cls, model_names=['rf', 'xgb', 'lgb', 'logistic']
        )
        
        print("模型训练完成")
    
    def _evaluate_models(self):
        """模型评估步骤"""
        # 加载特征数据
        features_file = os.path.join(self.data_dir, "features", "gold_features.csv")
        features_data = pd.read_csv(features_file)
        
        # 评估回归模型
        print("评估回归模型...")
        X_train, X_test, y_train, y_test, feature_cols = self.ml_models.prepare_data(
            features_data, target_col='target_return_1d', test_size=0.2
        )
        
        regression_results = self.ml_models.evaluate_regression_models(X_test, y_test)
        self.ml_models.save_results(regression_results, 'regression_results.csv')
        
        print("回归模型评估结果:")
        print(regression_results)
        
        # 评估分类模型
        print("\n评估分类模型...")
        X_train_cls, X_test_cls, y_train_cls, y_test_cls, feature_cols_cls = self.ml_models.prepare_data(
            features_data, target_col='target_direction_1d', test_size=0.2
        )
        
        classification_results = self.ml_models.evaluate_classification_models(X_test_cls, y_test_cls)
        self.ml_models.save_results(classification_results, 'classification_results.csv')
        
        print("\n分类模型评估结果:")
        print(classification_results)
    
    def predict_latest(self, model_name='ridge', model_type='regression'):
        """
        使用训练好的模型预测最新数据
        
        Args:
            model_name (str): 模型名称
            model_type (str): 模型类型 ('regression' 或 'classification')
            
        Returns:
            dict: 预测结果
        """
        print(f"使用 {model_name} 模型进行预测...")
        
        # 加载特征数据
        features_file = os.path.join(self.data_dir, "features", "gold_features.csv")
        features_data = pd.read_csv(features_file)
        
        # 准备数据
        target_col = 'target_return_1d' if model_type == 'regression' else 'target_direction_1d'
        X_train, X_test, y_train, y_test, feature_cols = self.ml_models.prepare_data(
            features_data, target_col=target_col, test_size=0.2
        )
        
        # 加载模型
        model = self.ml_models.load_model(model_name, model_type)
        if model is None:
            return {"error": f"无法加载模型 {model_name}"}
        
        # 获取最新数据（最后一行）
        latest_data = X_test.iloc[-1:].copy()
        
        # 进行预测
        prediction = model.predict(latest_data)[0]
        
        # 获取实际值（如果有）
        actual = y_test.iloc[-1] if len(y_test) > 0 else None
        
        result = {
            "model_name": model_name,
            "model_type": model_type,
            "prediction": prediction,
            "actual": actual,
            "date": features_data['date'].iloc[-1],
            "latest_price": features_data['close'].iloc[-1] if 'close' in features_data.columns else None
        }
        
        if model_type == 'regression':
            result["predicted_return"] = prediction
            result["predicted_price"] = result["latest_price"] * (1 + prediction) if result["latest_price"] else None
        else:
            result["predicted_direction"] = "上涨" if prediction > 0.5 else "下跌"
        
        print(f"预测完成: {result}")
        return result
    
    def get_model_summary(self):
        """获取模型训练摘要"""
        print("模型训练摘要:")
        print("-" * 40)
        
        # 检查模型文件
        model_files = []
        for file in os.listdir(self.model_dir):
            if file.endswith('.joblib'):
                model_files.append(file)
        
        print(f"已训练模型数量: {len(model_files)}")
        for file in model_files:
            print(f"  - {file}")
        
        # 检查结果文件
        result_files = []
        for file in os.listdir(self.model_dir):
            if file.endswith('.csv'):
                result_files.append(file)
        
        print(f"评估结果文件数量: {len(result_files)}")
        for file in result_files:
            print(f"  - {file}")
        
        # 显示最佳模型
        regression_results_file = os.path.join(self.model_dir, 'regression_results.csv')
        if os.path.exists(regression_results_file):
            regression_results = pd.read_csv(regression_results_file)
            if not regression_results.empty:
                best_reg = regression_results.iloc[0]
                print(f"\n最佳回归模型: {best_reg['Model']}")
                print(f"  RMSE: {best_reg['RMSE']:.6f}")
                print(f"  方向准确率: {best_reg['Direction_Accuracy']:.4f}")
        
        classification_results_file = os.path.join(self.model_dir, 'classification_results.csv')
        if os.path.exists(classification_results_file):
            classification_results = pd.read_csv(classification_results_file)
            if not classification_results.empty:
                best_cls = classification_results.iloc[0]
                print(f"\n最佳分类模型: {best_cls['Model']}")
                print(f"  准确率: {best_cls['Accuracy']:.4f}")
                print(f"  F1分数: {best_cls['F1']:.4f}")


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='黄金交易预测系统')
    parser.add_argument('--mode', choices=['full', 'predict', 'summary'], 
                       default='full', help='运行模式')
    parser.add_argument('--model', default='ridge', help='模型名称')
    parser.add_argument('--model_type', choices=['regression', 'classification'], 
                       default='regression', help='模型类型')
    parser.add_argument('--start_date', default='2014-01-01', help='开始日期')
    parser.add_argument('--end_date', default=None, help='结束日期')
    
    args = parser.parse_args()
    
    # 创建预测器
    predictor = GoldTradingPredictor()
    
    if args.mode == 'full':
        # 运行完整流水线
        predictor.run_full_pipeline(args.start_date, args.end_date)
        
    elif args.mode == 'predict':
        # 进行预测
        result = predictor.predict_latest(args.model, args.model_type)
        print("\n预测结果:")
        for key, value in result.items():
            print(f"{key}: {value}")
            
    elif args.mode == 'summary':
        # 显示模型摘要
        predictor.get_model_summary()


if __name__ == "__main__":
    # 如果没有命令行参数，使用默认模式
    if len(sys.argv) == 1:
        print("运行完整流水线...")
        predictor = GoldTradingPredictor()
        predictor.run_full_pipeline()
    else:
        main()
